#!/bin/bash
#############################################################################
#  Copyright (C) 2013-2015 Lawrence Livermore National Security, LLC.
#  Produced at Lawrence Livermore National Laboratory (cf, DISCLAIMER).
#  Written by Albert Chu <chu11@llnl.gov>
#  LLNL-CODE-644248
#
#  This file is part of Magpie, scripts for running Hadoop on
#  traditional HPC systems.  For details, see https://github.com/llnl/magpie.
#
#  Magpie is free software; you can redistribute it and/or modify it
#  under the terms of the GNU General Public License as published by
#  the Free Software Foundation; either version 2 of the License, or
#  (at your option) any later version.
#
#  Magpie is distributed in the hope that it will be useful, but
#  WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
#  General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with Magpie.  If not, see <http://www.gnu.org/licenses/>.
#############################################################################

# These are functions to be called by magpie-run

source ${MAGPIE_SCRIPTS_HOME}/magpie/exports/magpie-exports-submission-type
source ${MAGPIE_SCRIPTS_HOME}/magpie/exports/magpie-exports-dirs
source ${MAGPIE_SCRIPTS_HOME}/magpie/exports/magpie-exports-user
source ${MAGPIE_SCRIPTS_HOME}/magpie/lib/magpie-lib-defaults
source ${MAGPIE_SCRIPTS_HOME}/magpie/lib/magpie-lib-node-identification
source ${MAGPIE_SCRIPTS_HOME}/magpie/lib/magpie-lib-run
source ${MAGPIE_SCRIPTS_HOME}/magpie/lib/magpie-lib-paths

Magpie_run_start_spark () {
    if [ "${SPARK_SETUP}" == "yes" ] && [ "${magpie_run_prior_startup_successful}" == "true" ]
    then
        cd ${SPARK_HOME}

        if [ ${MAGPIE_JOB_TYPE} != "setuponly" ]
        then
            if [ "${SPARK_SETUP_TYPE}" == "STANDALONE" ]
            then
                # Make variables unspecified for launching
                Magpie_make_all_local_dirs_unspecified

                echo "Starting spark"
                ${sparksetupscriptprefix}/start-all.sh

                # Make variables specific now within Magpie
                Magpie_make_all_local_dirs_node_specific

                # My rough estimate for setup time is 30 seconds per 128 nodes
                local sleepwait=`expr ${SPARK_WORKER_COUNT} \/ 128 \* 30`
                if [ ${sleepwait} -lt 30 ]
                then
                    sleepwait=30
                fi
                echo "Waiting ${sleepwait} seconds to allow Spark daemons to setup"
                sleep ${sleepwait}
                magpie_run_total_sleep_wait=`expr ${magpie_run_total_sleep_wait} + ${sleepwait}`
            fi
        fi

        echo "*******************************************************"
        echo "*"
        echo "* Spark Information"
        echo "*"
        if [ "${SPARK_SETUP_TYPE}" == "STANDALONE" ]
        then
            echo "* You can view your Spark status by launching a web browser and pointing to ..."
            echo "*"
            echo "* Spark Master: http://${SPARK_MASTER_NODE}:${default_spark_master_webui_port}"
            echo "* Spark Worker: http://<WORKERNODE>:${default_spark_worker_webui_port}"
            echo "* Spark Application Dashboard: http://${SPARK_MASTER_NODE}:${default_spark_application_dashboard_port}"
            echo "*"
            echo "* The Spark Master for running jobs is"
            echo "*"
            echo "* spark://${SPARK_MASTER_NODE}:${default_spark_master_port}"
        fi
        echo "*"
        echo "* To access Spark directly, you'll want to:"
        echo "*"
        echo "*   ${MAGPIE_REMOTE_CMD:-ssh}${MAGPIE_REMOTE_CMD_OPTS:+" "}${MAGPIE_REMOTE_CMD_OPTS} ${SPARK_MASTER_NODE}"
        if echo $MAGPIE_SHELL | grep -q csh
        then
            echo "*   setenv JAVA_HOME \"${JAVA_HOME}\""
            echo "*   setenv SPARK_HOME \"${SPARK_HOME}\""
            echo "*   setenv SPARK_CONF_DIR \"${SPARK_CONF_DIR}\""
        else
            echo "*   export JAVA_HOME=\"${JAVA_HOME}\""
            echo "*   export SPARK_HOME=\"${SPARK_HOME}\""
            echo "*   export SPARK_CONF_DIR=\"${SPARK_CONF_DIR}\""
        fi
        echo "*"
        echo "* Then you can do as you please.  For example to run a job:"
        echo "*"
        echo "*   \$SPARK_HOME/${sparkcmdprefix}/spark-submit --class <class> <jar>"
        echo "*"

        if [ "${MAGPIE_JOB_TYPE}" == "setuponly" ]
        then
            echo "* To setup, login and set environment variables per the"
            echo "* instructions above, then run:"
            echo "*"
            echo "*   \$SPARK_HOME/${sparksetupscriptprefix}/start-all.sh"
            echo "*"
            echo "* To end/cleanup your session & kill all daemons, login and set"
            echo "* environment variables per the instructions above, then run:"
            echo "*"
            echo "*   \$SPARK_HOME/${sparksetupscriptprefix}/stop-all.sh"
            echo "*"
        fi
        if [ "${MAGPIE_ENVIRONMENT_VARIABLE_SCRIPT}X" != "X" ]
        then
            echo "* If running interactively, sourcing"
            echo "*"
            echo "* ${MAGPIE_ENVIRONMENT_VARIABLE_SCRIPT}"
            echo "*"
            echo "* will set most common environment variables for your job."
            echo "*"
        fi
        echo "*******************************************************"

        if [ "${MAGPIE_ENVIRONMENT_VARIABLE_SCRIPT}X" != "X" ]
        then
            if echo $MAGPIE_SHELL | grep -q csh
            then
                echo "setenv SPARK_HOME \"${SPARK_HOME}\"" >> ${MAGPIE_ENVIRONMENT_VARIABLE_SCRIPT}
                echo "setenv SPARK_CONF_DIR \"${SPARK_CONF_DIR}\"" >> ${MAGPIE_ENVIRONMENT_VARIABLE_SCRIPT}
                echo "setenv SPARK_LOG_DIR \"${SPARK_LOG_DIR}\"" >> ${MAGPIE_ENVIRONMENT_VARIABLE_SCRIPT}
                echo "setenv SPARK_MASTER_NODE \"${SPARK_MASTER_NODE}\"" >> ${MAGPIE_ENVIRONMENT_VARIABLE_SCRIPT}
                if [ "${SPARK_SETUP_TYPE}" == "STANDALONE" ]
                then
                    echo "setenv SPARK_MASTER_PORT \"${default_spark_master_port}\"" >> ${MAGPIE_ENVIRONMENT_VARIABLE_SCRIPT}
                fi
                echo "setenv SPARK_WORKER_COUNT \"${SPARK_WORKER_COUNT}\"" >> ${MAGPIE_ENVIRONMENT_VARIABLE_SCRIPT}
                echo "setenv SPARK_WORKER_CORE_COUNT \"${SPARK_WORKER_CORE_COUNT}\"" >> ${MAGPIE_ENVIRONMENT_VARIABLE_SCRIPT}
                echo "setenv SPARK_VERSION \"${SPARK_VERSION}\"" >> ${MAGPIE_ENVIRONMENT_VARIABLE_SCRIPT}
            else
                echo "export SPARK_HOME=\"${SPARK_HOME}\"" >> ${MAGPIE_ENVIRONMENT_VARIABLE_SCRIPT}
                echo "export SPARK_CONF_DIR=\"${SPARK_CONF_DIR}\"" >> ${MAGPIE_ENVIRONMENT_VARIABLE_SCRIPT}
                echo "export SPARK_LOG_DIR=\"${SPARK_LOG_DIR}\"" >> ${MAGPIE_ENVIRONMENT_VARIABLE_SCRIPT}
                echo "export SPARK_MASTER_NODE=\"${SPARK_MASTER_NODE}\"" >> ${MAGPIE_ENVIRONMENT_VARIABLE_SCRIPT}
                if [ "${SPARK_SETUP_TYPE}" == "STANDALONE" ]
                then
                    echo "export SPARK_MASTER_PORT=\"${default_spark_master_port}\"" >> ${MAGPIE_ENVIRONMENT_VARIABLE_SCRIPT}
                fi
                echo "export SPARK_WORKER_COUNT=\"${SPARK_WORKER_COUNT}\"" >> ${MAGPIE_ENVIRONMENT_VARIABLE_SCRIPT}
                echo "export SPARK_WORKER_CORE_COUNT=\"${SPARK_WORKER_CORE_COUNT}\"" >> ${MAGPIE_ENVIRONMENT_VARIABLE_SCRIPT}
                echo "export SPARK_VERSION=\"${SPARK_VERSION}\"" >> ${MAGPIE_ENVIRONMENT_VARIABLE_SCRIPT}
            fi
            echo "" >> ${MAGPIE_ENVIRONMENT_VARIABLE_SCRIPT}
        fi

        # Nothing to check for Spark, setup always passes
        magpie_run_spark_should_be_torndown=1
        magpie_run_spark_setup_successful=1
    else
        magpie_run_spark_should_be_torndown=0
        magpie_run_spark_setup_successful=1
    fi
}

Magpie_run_spark () {
    if [ "${SPARK_JOB}" == "sparkpi" ]
    then
        echo "*******************************************************"
        echo "* Running SparkPi"
        echo "*******************************************************"
        ${MAGPIE_SCRIPTS_HOME}/magpie/run/magpie-run-execute script ${MAGPIE_SCRIPTS_HOME}/magpie/job/magpie-job-spark-sparkpi &
        local scriptpid=$!
        Magpie_wait_script_sigusr2_on_job_timeout ${scriptpid}
    elif [ "${SPARK_JOB}" == "sparkwordcount" ]
    then
        echo "*******************************************************"
        echo "* Running SparkWordCount"
        echo "*******************************************************"
        ${MAGPIE_SCRIPTS_HOME}/magpie/run/magpie-run-execute script ${MAGPIE_SCRIPTS_HOME}/magpie/job/magpie-job-spark-sparkwordcount &
        local scriptpid=$!
        Magpie_wait_script_sigusr2_on_job_timeout ${scriptpid}
    else
        Magpie_output_internal_error "SPARK_JOB = ${SPARK_JOB} not handled"
    fi
}

Magpie_run_stop_spark () {
    if [ "${SPARK_SETUP}" == "yes" ] && [ "${magpie_run_spark_should_be_torndown}" == "1" ]
    then
        if [ ${MAGPIE_JOB_TYPE} != "setuponly" ]
        then
            if [ "${SPARK_SETUP_TYPE}" == "STANDALONE" ]
            then
                cd ${SPARK_HOME}

                # Make variables unspecified for shutdown
                Magpie_make_all_local_dirs_unspecified

                echo "Stopping spark"
                ${sparksetupscriptprefix}/stop-all.sh

                # Make variables specific now within Magpie
                Magpie_make_all_local_dirs_node_specific
            fi
        fi
    fi
    magpie_run_spark_teardown_complete=1
}
